
from matplotlib import pyplot as plt
from torch.utils.data import DataLoader
from torchvision import transforms

from main.utils.pathmanager import data_path

labelToStr = (
    '飞机 (airplane)', '汽车 (automobile)', '鸟类 (bird)', '猫 (cat)', '鹿 (deer)', '狗 (dog)', '蛙 (frog)',
    '马 (horse)', '船 (ship)', '卡车 (truck)')


def data_set(dataset, batch_size=128, num_workers=4):
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
    ])

    def _init(train, shuffle):
        _data_path = data_path()
        kwargs = dict(root=_data_path,
                      train=train,
                      download=False,
                      transform=transform)
        _dataset = dataset(**kwargs)
        return DataLoader(_dataset, batch_size=batch_size,
                          shuffle=shuffle,
                          num_workers=num_workers,
                          pin_memory=True,  # 提高数据加载速度
                          drop_last=True  # 丢弃最后一个不足batch的数据
                          )
    return _init


def look_look(loader):
    # 获取随机训练集的图片
    dataiter = iter(loader)
    # 获取一个批次的数据
    images, labels = next(dataiter)
    show_img(img=images[0], label=labels[0])


def check_out(img, label, sub):
    show_img(img, label, lambda p: p.suptitle(sub))


def show_img(img, label, func=None):
    # 指定字体路径，如果字体已经安装在系统中，可以直接使用字体名，如 'SimHei'
    plt.rcParams['font.sans-serif'] = ['SimHei']
    image = img.numpy().transpose((1, 2, 0))
    # 显示图片,bilinear,nearest,bicubic
    # plt.imshow(image, interpolation='bicubic')
    plt.imshow(image)
    if func is not None:
        func(plt)
    plt.title(f'Label: {labelToStr[label]}')
    # plt.title(labelToStr[label])
    plt.show()


if __name__ == '__main__':
    from torchvision import datasets
    cifar10 = data_set(datasets.CIFAR10)
    trainloader = cifar10(True, True)
    testloader = cifar10(False, False)
    look_look(testloader)
